// File: jmcomic-core/src/main/java/dev/jukomu/common/base/concurrent/StopFlag.java
package dev.jukomu.common.base.concurrent;

import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 对应Python的StopThreadFlag，用于管理线程的停止信号。
 * 允许标记单个线程或所有线程停止。
 */
public class StopFlag {

    public static final boolean RUN = true;
    public static final boolean STOP = false;

    private final String key; // 标记线程属性的键名，例如 "shouldStop"
    private final Set<Thread> markedThreadSet; // 跟踪被此Flag标记过的线程
    private final AtomicBoolean allStop; // 标记所有线程是否应该停止

    /**
     * 构造函数。
     *
     * @param key 用于标记线程的属性键。
     */
    public StopFlag(String key) {
        this.key = key;
        this.markedThreadSet = ConcurrentHashMap.newKeySet(); // 线程安全的Set
        this.allStop = new AtomicBoolean(false);
    }

    /**
     * 检查给定线程是否应该停止。
     * 如果线程未被明确标记为运行，则默认标记为运行。
     *
     * @param thread 要检查的线程，如果为null则使用当前线程。
     * @return 如果线程应该停止返回 true。
     */
    public boolean shouldStop(Thread thread) {
        thread = (thread != null) ? thread : Thread.currentThread();

        if (allStop.get()) {
            return true;
        }

        // 检查线程局部变量
        // Java中没有直接的setattr/getattr，我们通过ThreadLocal或其他方式模拟
        // 为简化，这里直接检查线程的中断状态，并假设外部协作设置停止标志
        // 或者，我们可以将一个Map<Thread, Boolean>作为内部状态来模拟getattr/setattr
        // 为了模拟Python的getattr/setattr，我们将使用一个ConcurrentHashMap来保存每个线程的停止状态。
        return !markedThreadSet.contains(thread) || markedThreadSet.contains(thread) && (Boolean) getThreadLocalFlag(thread) == STOP;
    }

    /**
     * 检查当前线程是否应该停止。
     *
     * @return 如果当前线程应该停止返回 true。
     */
    public boolean shouldStop() {
        return shouldStop(Thread.currentThread());
    }

    /**
     * 标记线程为运行状态。
     * 如果线程已经被标记为停止或全部停止，则打印警告。
     *
     * @param thread 要标记的线程，如果为null则使用当前线程。
     */
    public void markRun(Thread thread) {
        thread = (thread != null) ? thread : Thread.currentThread();

        if (allStop.get() || (markedThreadSet.contains(thread) && (Boolean) getThreadLocalFlag(thread) == STOP)) {
            System.err.println(
                    key + " is already set STOP for thread: [" + thread.getName() + "], " +
                            "current thread: [" + Thread.currentThread().getName() + "]"
            );
            return;
        }

        markedThreadSet.add(thread);
        setThreadLocalFlag(thread, RUN);
    }

    /**
     * 标记当前线程为运行状态。
     */
    public void markRun() {
        markRun(Thread.currentThread());
    }

    /**
     * 标记指定线程为停止状态。
     *
     * @param thread 要标记的线程，如果为null则使用当前线程。
     */
    public void markStop(Thread thread) {
        thread = (thread != null) ? thread : Thread.currentThread();
        markedThreadSet.add(thread); // 确保线程被跟踪
        setThreadLocalFlag(thread, STOP);
    }

    /**
     * 标记所有被跟踪的线程为停止状态。
     */
    public void markStopForAll() {
        allStop.set(true);
        for (Thread thread : markedThreadSet) {
            markStop(thread);
        }
    }

    /**
     * 获取所有被此Flag标记的线程集合。
     *
     * @return 线程集合。
     */
    public Set<Thread> getMarkedThreadSet() {
        return markedThreadSet;
    }

    /**
     * 等待所有被标记的线程完成。
     *
     * @throws InterruptedException 如果当前线程在等待期间被中断。
     */
    public void waitForAll() throws InterruptedException {
        for (Thread thread : markedThreadSet) {
            if (thread == Thread.currentThread()) {
                continue;
            }
            while (thread.isAlive()) {
                thread.join(100); // Wait up to 100ms, then recheck shouldStop
            }
        }
    }

    /**
     * 抛出停止异常（模拟KeyboardInterrupt）。
     *
     * @throws InterruptedException 模拟停止信号。
     */
    public void raiseStopException() throws InterruptedException {
        throw new InterruptedException("Stop signal received.");
    }

    /**
     * 休眠一段时间，并在休眠期间定期检查停止标志。
     * 如果应该停止，则抛出停止异常。
     *
     * @param durationSeconds 休眠总时长（秒）。
     * @throws InterruptedException 如果应该停止或休眠被中断。
     */
    public void sleepOrRaise(double durationSeconds) throws InterruptedException {
        if (durationSeconds < 0.1) {
            TimeUnit.MILLISECONDS.sleep((long) (durationSeconds * 1000));
            ifStopRaise();
            return;
        }

        long totalMillis = (long) (durationSeconds * 1000);
        long checkIntervalMillis = 100; // Check every 100ms

        long startTime = System.currentTimeMillis();
        while (System.currentTimeMillis() - startTime < totalMillis) {
            ifStopRaise();
            long remaining = totalMillis - (System.currentTimeMillis() - startTime);
            if (remaining <= 0) break;
            TimeUnit.MILLISECONDS.sleep(Math.min(checkIntervalMillis, remaining));
        }
        ifStopRaise(); // Final check after loop
    }

    /**
     * 如果应该停止，则抛出停止异常。
     *
     * @throws InterruptedException 如果应该停止。
     */
    public void ifStopRaise() throws InterruptedException {
        if (shouldStop()) {
            raiseStopException();
        }
    }

    /**
     * 休眠一段时间，并在休眠期间定期检查停止标志。
     *
     * @param durationSeconds 休眠总时长（秒）。
     * @param partSeconds     每次检查停止标志的间隔（秒）。
     * @return 如果在休眠期间收到停止信号返回 true，否则返回 false。
     * @throws InterruptedException 如果休眠被中断。
     */
    public boolean sleepOrReturn(double durationSeconds, double partSeconds) throws InterruptedException {
        if (durationSeconds < partSeconds) {
            TimeUnit.MILLISECONDS.sleep((long) (durationSeconds * 1000));
            return shouldStop();
        }

        long totalMillis = (long) (durationSeconds * 1000);
        long partMillis = (long) (partSeconds * 1000);

        long startTime = System.currentTimeMillis();
        while (System.currentTimeMillis() - startTime < totalMillis) {
            if (shouldStop()) {
                return true;
            }
            long remaining = totalMillis - (System.currentTimeMillis() - startTime);
            if (remaining <= 0) break;
            TimeUnit.MILLISECONDS.sleep(Math.min(partMillis, remaining));
        }
        return shouldStop(); // Final check
    }

    /**
     * 休眠一段时间，并在休眠期间定期检查停止标志。
     * 使用默认的检查间隔0.5秒。
     *
     * @param durationSeconds 休眠总时长（秒）。
     * @return 如果在休眠期间收到停止信号返回 true，否则返回 false。
     * @throws InterruptedException 如果休眠被中断。
     */
    public boolean sleepOrReturn(double durationSeconds) throws InterruptedException {
        return sleepOrReturn(durationSeconds, 0.5);
    }

    // --- 模拟Python的getattr/setattr来存储线程局部标志 ---
    // Python的getattr/setattr直接修改对象属性，Java没有。
    // 这里使用ConcurrentHashMap来模拟，将线程作为键，布尔值作为值。
    private final Map<Thread, Boolean> threadLocalFlags = new ConcurrentHashMap<>();

    private Boolean getThreadLocalFlag(Thread thread) {
        return threadLocalFlags.getOrDefault(thread, RUN); // 默认是RUN
    }

    private void setThreadLocalFlag(Thread thread, Boolean value) {
        threadLocalFlags.put(thread, value);
    }
}